"""
Author: Morphlng
Date: 2024-04-18 17:00:31
LastEditTime: 2024-04-23 16:44:39
LastEditors: Morphlng
Description: Abstract action interface for RL agent.
FilePath: /DrivingGym/src/driving_gym/environment/agent/action/action_interface.py
"""

from __future__ import annotations

import gymnasium as gym
import numpy as np

from driving_gym.simulation.adapter_interface import AdapterInterface


class AbstractAction:
    def __init__(self, action, duration: int):
        """An action can last for multiple steps and return a control signal dict at each step."""
        self.action = action
        self.duration = duration

    def run_step(self) -> dict:
        """Run the action for one step and return a control signal dict."""
        raise NotImplementedError

    def done(self) -> bool:
        """Signal whether the action is done."""
        return self.duration <= 0

    def to_dict(self) -> dict:
        """Return a dictionary representation of the action."""
        return {"action": self.action, "duration": self.duration}


class ActionInterface:
    def __init__(self, config: dict, adapter: AdapterInterface):
        """Initialize the action handler with a configuration dictionary."""
        self.action_config = config.copy()
        self.adapter = adapter

        # This is automatically set by the environment
        self.actor_id: str = self.action_config["actor_id"]

    def reset(self):
        """Reset the action handler."""
        pass

    def convert_single_action(self, action) -> AbstractAction:
        """Convert a single action to a AbstactAction object.

        Args:
            action: model output action

        Returns:
            AbstactAction: The converted action. Call run_step with to get the control signal.
        """
        raise NotImplementedError

    def get_action_space(self) -> gym.spaces.Space:
        """Get the action space for the agent."""
        raise NotImplementedError

    def get_action_mask(self) -> np.ndarray:
        """Check if all the action in the action space is valid for the actor.

        Returns:
            np.ndarray: An int8 array indicating the valid actions, None if no mask is needed.
        """
        return None

    def get_mask_space(self) -> gym.spaces.Space:
        """Get the mask space for the agent.

        Note:
            Override and return None if no mask space is needed.
        """

        action_space = self.get_action_space()
        if isinstance(action_space, gym.spaces.Discrete):
            return gym.spaces.Box(
                low=0,
                high=1,
                shape=(action_space.n,),
                dtype=np.int8,
            )
        elif isinstance(action_space, gym.spaces.MultiDiscrete):
            return gym.spaces.Box(
                low=0,
                high=1,
                shape=sum(action_space.nvec),
                dtype=np.int8,
            )
        else:
            raise ValueError(f"Unsupported action space type: {type(action_space)}")
